"""
NLP 自然语言处理任务一般流程
1. 获取数据 【从网页爬取数据，存入TXT csv 等】
2. 缓存到内存上 【加载数据】
3. 预处理 【清洗 分词】
4. 词嵌入 【目的将文本数据转换为计算机能理解的数值】
5. RNN/LSTM/GRU 直接训练
6. transformer 位置编码 注意力机制 掩码

"""
import pandas as pd
import json
from torch.utils.data import Dataset, DataLoader

class LanguageDataset(Dataset):
    def __init__(self, max_row=1000):
        super().__init__()
        # 加载数据  将csv文件数据缓存到内存
        self.df = pd.read_csv('../data/chatData.csv', nrows=max_row, encoding='utf-8')

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 按照下标索引idx获取其中一条数据  也就是一个问答对
        row = self.df.iloc[idx]
        # 提取问答对
        conversation = json.loads(row['conversation'])[0]
        # 提取human 问
        src = conversation['human']
        # 提取assistant 答
        assistant = conversation['assistant']
        # 设置开始标志CLS 和 结束标志SEP
        tgt = '[CLS]' + assistant
        label = assistant + '[SEP]'
        return src, tgt, label

class Dataset(Dataset):
    def __init__(self, max_row=1000):
        super().__init__()
        # 加载数据  将csv文件数据缓存到内存
        self.df = pd.read_csv('../data/chatData.csv', nrows=max_row, encoding='utf-8')

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # 按照下标索引idx获取其中一条数据  也就是一个问答对
        row = self.df.iloc[idx]
        # 提取问答对
        conversation = json.loads(row['conversation'])[0]
        # 提取human 问
        src = conversation['human']
        # 提取assistant 答
        assistant = conversation['assistant']
        # 设置开始标志CLS 和 结束标志SEP
        tgt = '[CLS]' + assistant
        label = assistant + '[SEP]'
        return src, tgt, label

if __name__ == '__main__':
    ds = LanguageDataset(max_row=500)
    print(ds[0])
    # 按批次加载数据
    dl = DataLoader(ds, batch_size=5)
    for src, tgt, label in dl:
        print(src, tgt, label)









